{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先，需要了解庞加莱球模型中的指数映射（Exponential Map）函数，这是从欧氏空间到庞加莱球模型的关键映射。庞加莱球模型的指数映射定义如下：\n",
    "\n",
    "$$\n",
    "\\text{exp}_c^x(v) = x \\oplus_c \\left( \\tanh\\left( \\sqrt{c} \\frac{\\lambda_x \\|v\\|}{2} \\right) \\frac{v}{\\sqrt{c} \\|v\\|} \\right)\n",
    "$$\n",
    "\n",
    "其中，$\\oplus_c$ 是Möbius加法，定义如下：\n",
    "\n",
    "$$\n",
    "x \\oplus_c y = \\frac{(1 + 2c \\langle x, y \\rangle + c \\|y\\|^2)x + (1 - c \\|x\\|^2)y}{1 + 2c \\langle x, y \\rangle + c^2 \\|x\\|^2 \\|y\\|^2}\n",
    "$$\n",
    "\n",
    "以及\n",
    "\n",
    "$$\n",
    "\\lambda_x = \\frac{2}{1 - c \\|x\\|^2}\n",
    "$$\n",
    "\n",
    "\n",
    "注意：\n",
    "\n",
    "1. 庞加莱球模型的基点`x`通常被设为原点（即全零向量），但在某些情况下，为了保持数值稳定性，可能会选择其他点。\n",
    "2. 在实际应用中，你可能需要处理批量数据，因此上述代码中的向量都是按批次（batch）组织的，例如`x`和`v`的形状都是`(batch_size, feature_dim)`。\n",
    "3. 曲率`c`是一个重要的超参数，其值会影响庞加莱球的大小和形状，需要根据具体任务调整。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在PyTorch中实现这个转换，我们首先需要定义Möbius加法和指数映射函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "sqrt(): argument 'input' (position 1) must be Tensor, not float",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m/root/MyCode/LinearSpace/Poincaré_ball_model.ipynb Cell 3\u001b[0m line \u001b[0;36m3\n\u001b[1;32m     <a href='vscode-notebook-cell://localhost:8080/root/MyCode/LinearSpace/Poincar%C3%A9_ball_model.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=34'>35</a>\u001b[0m v \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mrandn(\u001b[39m1\u001b[39m, \u001b[39m5\u001b[39m)  \u001b[39m# 从x出发的向量\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://localhost:8080/root/MyCode/LinearSpace/Poincar%C3%A9_ball_model.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=36'>37</a>\u001b[0m \u001b[39m# 映射到庞加莱球\u001b[39;00m\n\u001b[0;32m---> <a href='vscode-notebook-cell://localhost:8080/root/MyCode/LinearSpace/Poincar%C3%A9_ball_model.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=37'>38</a>\u001b[0m y \u001b[39m=\u001b[39m exp_map(x, v, c)\n\u001b[1;32m     <a href='vscode-notebook-cell://localhost:8080/root/MyCode/LinearSpace/Poincar%C3%A9_ball_model.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=38'>39</a>\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mMapped point in Poincaré ball:\u001b[39m\u001b[39m\"\u001b[39m, y)\n",
      "\u001b[1;32m/root/MyCode/LinearSpace/Poincaré_ball_model.ipynb Cell 3\u001b[0m line \u001b[0;36m2\n\u001b[1;32m     <a href='vscode-notebook-cell://localhost:8080/root/MyCode/LinearSpace/Poincar%C3%A9_ball_model.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25'>26</a>\u001b[0m v_norm \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mnorm(v, dim\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m, keepdim\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m)\n\u001b[1;32m     <a href='vscode-notebook-cell://localhost:8080/root/MyCode/LinearSpace/Poincar%C3%A9_ball_model.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=26'>27</a>\u001b[0m lambda_x \u001b[39m=\u001b[39m \u001b[39m2\u001b[39m \u001b[39m/\u001b[39m (\u001b[39m1\u001b[39m \u001b[39m-\u001b[39m c \u001b[39m*\u001b[39m torch\u001b[39m.\u001b[39msum(x \u001b[39m*\u001b[39m\u001b[39m*\u001b[39m \u001b[39m2\u001b[39m, dim\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m, keepdim\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m))\n\u001b[0;32m---> <a href='vscode-notebook-cell://localhost:8080/root/MyCode/LinearSpace/Poincar%C3%A9_ball_model.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=27'>28</a>\u001b[0m scaled_v \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mtanh(torch\u001b[39m.\u001b[39;49msqrt(c) \u001b[39m*\u001b[39m lambda_x \u001b[39m/\u001b[39m \u001b[39m2\u001b[39m \u001b[39m*\u001b[39m v_norm) \u001b[39m*\u001b[39m v \u001b[39m/\u001b[39m torch\u001b[39m.\u001b[39msqrt(c) \u001b[39m/\u001b[39m v_norm\n\u001b[1;32m     <a href='vscode-notebook-cell://localhost:8080/root/MyCode/LinearSpace/Poincar%C3%A9_ball_model.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=28'>29</a>\u001b[0m \u001b[39mreturn\u001b[39;00m mobius_add(x, scaled_v, c)\n",
      "\u001b[0;31mTypeError\u001b[0m: sqrt(): argument 'input' (position 1) must be Tensor, not float"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "def mobius_add(x, y, c):\n",
    "    \"\"\"\n",
    "    Möbius 加法\n",
    "    :param x: torch.Tensor, 点x的坐标\n",
    "    :param y: torch.Tensor, 点y的坐标\n",
    "    :param c: float, 庞加莱球的曲率\n",
    "    :return: torch.Tensor, x和y的Möbius和\n",
    "    \"\"\"\n",
    "    x2 = torch.sum(x ** 2, dim=1, keepdim=True)\n",
    "    y2 = torch.sum(y ** 2, dim=1, keepdim=True)\n",
    "    xy = torch.mm(x, y.t())\n",
    "    numerator = (1 + 2 * c * xy + c * y2) * x + (1 - c * x2) * y\n",
    "    denominator = 1 + 2 * c * xy + c ** 2 * x2 * y2\n",
    "    return numerator / denominator\n",
    "\n",
    "def exp_map(x, v, c):\n",
    "    \"\"\"\n",
    "    庞加莱球模型的指数映射\n",
    "    :param x: torch.Tensor, 基点x的坐标\n",
    "    :param v: torch.Tensor, 从x出发的向量\n",
    "    :param c: float, 庞加莱球的曲率\n",
    "    :return: torch.Tensor, 映射后的坐标\n",
    "    \"\"\"\n",
    "    v_norm = torch.norm(v, dim=1, keepdim=True)\n",
    "    lambda_x = 2 / (1 - c * torch.sum(x ** 2, dim=1, keepdim=True))\n",
    "    scaled_v = torch.tanh(torch.sqrt(c) * lambda_x / 2 * v_norm) * v / torch.sqrt(c) / v_norm\n",
    "    return mobius_add(x, scaled_v, c)\n",
    "\n",
    "# 示例\n",
    "c = 1.0  # 曲率，可根据需要调整\n",
    "x = torch.randn(1, 5)  # 基点x，随机生成\n",
    "x = x / torch.norm(x, dim=1, keepdim=True)  # 归一化到单位球内\n",
    "v = torch.randn(1, 5)  # 从x出发的向量\n",
    "\n",
    "# 映射到庞加莱球\n",
    "y = exp_map(x, v, c)\n",
    "print(\"Mapped point in Poincaré ball:\", y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "结合深度学习下的非欧空间"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch  \n",
    "import torch.nn as nn  \n",
    "import torch.nn.functional as F  \n",
    "  \n",
    "def hyperbolic_distance(x, y, c=1.0):  \n",
    "    \"\"\"  \n",
    "    计算两个点在Poincaré球模型中的双曲距离。  \n",
    "    x, y: PyTorch张量，表示双曲空间中的点。  \n",
    "    c: 双曲空间的曲率，默认为1.0。  \n",
    "    \"\"\"  \n",
    "    norm_x = torch.norm(x, p=2, dim=-1, keepdim=True)  \n",
    "    norm_y = torch.norm(y, p=2, dim=-1, keepdim=True)  \n",
    "      \n",
    "    # 使用Möbius加法的形式来避免直接的除法运算，这里我们直接计算双曲距离公式  \n",
    "    numerator = 1 + 2 * c * torch.sum(x * y, dim=-1, keepdim=True) + c * torch.norm(y, p=2, dim=-1, keepdim=True)**2  \n",
    "    denominator = 1 - c * torch.norm(x, p=2, dim=-1, keepdim=True)**2  \n",
    "    denominator *= 1 - c * torch.norm(y, p=2, dim=-1, keepdim=True)**2  \n",
    "      \n",
    "    distance = torch.acosh(numerator / denominator)  \n",
    "    return distance  \n",
    "  \n",
    "class HyperbolicLoss(nn.Module):  \n",
    "    def __init__(self, c=1.0):  \n",
    "        super(HyperbolicLoss, self).__init__()  \n",
    "        self.c = c  \n",
    "  \n",
    "    def forward(self, text_embeddings, image_embeddings):  \n",
    "        \"\"\"  \n",
    "        计算文本嵌入和图像嵌入在双曲空间中的平均距离作为损失。  \n",
    "        text_embeddings, image_embeddings: PyTorch张量，分别表示文本和图像的嵌入。  \n",
    "        \"\"\"  \n",
    "        distances = hyperbolic_distance(text_embeddings, image_embeddings, self.c)  \n",
    "        loss = torch.mean(distances)  \n",
    "        return loss  \n",
    "  \n",
    "# 假设我们有一些文本和图像的嵌入  \n",
    "# text_embeddings = torch.randn(batch_size, embedding_dim)  \n",
    "# image_embeddings = torch.randn(batch_size, embedding_dim)  \n",
    "  \n",
    "# 实例化损失函数并计算损失  \n",
    "# loss_fn = HyperbolicLoss()  \n",
    "# loss = loss_fn(text_embeddings, image_embeddings)  \n",
    "# print(loss)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "science39",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
